{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c9e82a03",
   "metadata": {},
   "source": [
    "# Transformer on MIMIC-IV (Demo)\n",
    "This minimal walkthrough uses PyHealth's built-in utilities to train the PyHealth Transformer model on a MIMIC-IV mortality prediction task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b3e69e36",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/logic/miniforge3/envs/pyhealth/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from pyhealth.datasets import MIMIC4Dataset, get_dataloader, split_by_sample\n",
    "from pyhealth.tasks import InHospitalMortalityMIMIC4\n",
    "from pyhealth.models import Transformer\n",
    "from pyhealth.trainer import Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7c246dee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage Starting MIMIC4Dataset init: 921.5 MB\n",
      "Initializing MIMIC4EHRDataset with tables: ['diagnoses_icd', 'procedures_icd', 'prescriptions', 'labevents'] (dev mode: True)\n",
      "Using default EHR config: /home/logic/miniforge3/envs/pyhealth/lib/python3.12/site-packages/pyhealth/datasets/configs/mimic4_ehr.yaml\n",
      "Memory usage Before initializing mimic4_ehr: 921.5 MB\n",
      "Initializing mimic4_ehr dataset from /home/logic/Github/mimic4 (dev mode: False)\n",
      "Scanning table: diagnoses_icd from /home/logic/Github/mimic4/hosp/diagnoses_icd.csv.gz\n",
      "Initializing MIMIC4EHRDataset with tables: ['diagnoses_icd', 'procedures_icd', 'prescriptions', 'labevents'] (dev mode: True)\n",
      "Using default EHR config: /home/logic/miniforge3/envs/pyhealth/lib/python3.12/site-packages/pyhealth/datasets/configs/mimic4_ehr.yaml\n",
      "Memory usage Before initializing mimic4_ehr: 921.5 MB\n",
      "Initializing mimic4_ehr dataset from /home/logic/Github/mimic4 (dev mode: False)\n",
      "Scanning table: diagnoses_icd from /home/logic/Github/mimic4/hosp/diagnoses_icd.csv.gz\n",
      "Joining with table: /home/logic/Github/mimic4/hosp/admissions.csv.gz\n",
      "Scanning table: procedures_icd from /home/logic/Github/mimic4/hosp/procedures_icd.csv.gz\n",
      "Joining with table: /home/logic/Github/mimic4/hosp/admissions.csv.gz\n",
      "Scanning table: prescriptions from /home/logic/Github/mimic4/hosp/prescriptions.csv.gz\n",
      "Scanning table: labevents from /home/logic/Github/mimic4/hosp/labevents.csv.gz\n",
      "Joining with table: /home/logic/Github/mimic4/hosp/admissions.csv.gz\n",
      "Scanning table: procedures_icd from /home/logic/Github/mimic4/hosp/procedures_icd.csv.gz\n",
      "Joining with table: /home/logic/Github/mimic4/hosp/admissions.csv.gz\n",
      "Scanning table: prescriptions from /home/logic/Github/mimic4/hosp/prescriptions.csv.gz\n",
      "Scanning table: labevents from /home/logic/Github/mimic4/hosp/labevents.csv.gz\n",
      "Joining with table: /home/logic/Github/mimic4/hosp/d_labitems.csv.gz\n",
      "Scanning table: patients from /home/logic/Github/mimic4/hosp/patients.csv.gz\n",
      "Scanning table: admissions from /home/logic/Github/mimic4/hosp/admissions.csv.gz\n",
      "Scanning table: icustays from /home/logic/Github/mimic4/icu/icustays.csv.gz\n",
      "Memory usage After initializing mimic4_ehr: 957.0 MB\n",
      "Memory usage After EHR dataset initialization: 957.0 MB\n",
      "Memory usage Before combining data: 957.0 MB\n",
      "Combining data from ehr dataset\n",
      "Creating combined dataframe\n",
      "Memory usage After combining data: 957.0 MB\n",
      "Memory usage Completed MIMIC4Dataset init: 957.0 MB\n",
      "Joining with table: /home/logic/Github/mimic4/hosp/d_labitems.csv.gz\n",
      "Scanning table: patients from /home/logic/Github/mimic4/hosp/patients.csv.gz\n",
      "Scanning table: admissions from /home/logic/Github/mimic4/hosp/admissions.csv.gz\n",
      "Scanning table: icustays from /home/logic/Github/mimic4/icu/icustays.csv.gz\n",
      "Memory usage After initializing mimic4_ehr: 957.0 MB\n",
      "Memory usage After EHR dataset initialization: 957.0 MB\n",
      "Memory usage Before combining data: 957.0 MB\n",
      "Combining data from ehr dataset\n",
      "Creating combined dataframe\n",
      "Memory usage After combining data: 957.0 MB\n",
      "Memory usage Completed MIMIC4Dataset init: 957.0 MB\n"
     ]
    }
   ],
   "source": [
    "dataset = MIMIC4Dataset(\n",
    "    ehr_root=\"/home/logic/Github/mimic4\",\n",
    "    ehr_tables=[\"diagnoses_icd\", \"procedures_icd\", \"prescriptions\", \"labevents\"],\n",
    "    dev=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "46bfc544",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setting task InHospitalMortalityMIMIC4 for mimic4 base dataset...\n",
      "Generating samples with 1 worker(s)...\n",
      "Collecting global event dataframe...\n",
      "Dev mode enabled: limiting to 1000 patients\n",
      "Generating samples with 1 worker(s)...\n",
      "Collecting global event dataframe...\n",
      "Dev mode enabled: limiting to 1000 patients\n",
      "Collected dataframe with shape: (131557, 47)\n",
      "Collected dataframe with shape: (131557, 47)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating samples for InHospitalMortalityMIMIC4 with 1 worker: 100%|██████████| 100/100 [00:01<00:00, 95.23it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Caching samples to ../../test_cache_transformer_m4/InHospitalMortalityMIMIC4.parquet\n",
      "Failed to cache samples: failed to determine supertype of list[datetime[μs]] and object\n",
      "Label mortality vocab: {0: 0, 1: 1}\n",
      "Failed to cache samples: failed to determine supertype of list[datetime[μs]] and object\n",
      "Label mortality vocab: {0: 0, 1: 1}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Processing samples: 100%|██████████| 216/216 [00:00<00:00, 1793.08it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated 216 samples for task InHospitalMortalityMIMIC4\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "task = InHospitalMortalityMIMIC4()\n",
    "sample_dataset = dataset.set_task(\n",
    "    task,\n",
    "    cache_dir=\"../../test_cache_transformer_m4\"\n",
    ")\n",
    "train_dataset, val_dataset, test_dataset = split_by_sample(sample_dataset, ratios=[0.7, 0.1, 0.2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9bc0c638",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = get_dataloader(train_dataset, batch_size=32, shuffle=True)\n",
    "val_loader = get_dataloader(val_dataset, batch_size=32, shuffle=False)\n",
    "test_loader = get_dataloader(test_dataset, batch_size=32, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5ea405d1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transformer(\n",
      "  (embedding_model): EmbeddingModel(embedding_layers=ModuleDict(\n",
      "    (labs): Linear(in_features=27, out_features=128, bias=True)\n",
      "  ))\n",
      "  (transformer): ModuleDict(\n",
      "    (labs): TransformerLayer(\n",
      "      (transformer): ModuleList(\n",
      "        (0-1): 2 x TransformerBlock(\n",
      "          (attention): MultiHeadedAttention(\n",
      "            (linear_layers): ModuleList(\n",
      "              (0-2): 3 x Linear(in_features=128, out_features=128, bias=False)\n",
      "            )\n",
      "            (output_linear): Linear(in_features=128, out_features=128, bias=False)\n",
      "            (attention): Attention()\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "          (feed_forward): PositionwiseFeedForward(\n",
      "            (w_1): Linear(in_features=128, out_features=512, bias=True)\n",
      "            (w_2): Linear(in_features=512, out_features=128, bias=True)\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "            (activation): GELU(approximate='none')\n",
      "          )\n",
      "          (input_sublayer): SublayerConnection(\n",
      "            (norm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "          (output_sublayer): SublayerConnection(\n",
      "            (norm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "          (dropout): Dropout(p=0.1, inplace=False)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (fc): Linear(in_features=128, out_features=1, bias=True)\n",
      ")\n",
      "Metrics: ['roc_auc', 'pr_auc']\n",
      "Device: cuda\n",
      "\n",
      "Metrics: ['roc_auc', 'pr_auc']\n",
      "Device: cuda\n",
      "\n"
     ]
    }
   ],
   "source": [
    "model = Transformer(dataset=sample_dataset, embedding_dim=128, heads=2, num_layers=2, dropout=0.1)\n",
    "trainer = Trainer(model=model, metrics=[\"roc_auc\", \"pr_auc\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "baf0ae3a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training:\n",
      "Batch size: 32\n",
      "Optimizer: <class 'torch.optim.adam.Adam'>\n",
      "Optimizer params: {'lr': 0.0001}\n",
      "Weight decay: 0.0\n",
      "Max grad norm: None\n",
      "Val dataloader: <torch.utils.data.dataloader.DataLoader object at 0x7f5e0a8b1fd0>\n",
      "Monitor: roc_auc\n",
      "Monitor criterion: max\n",
      "Epochs: 5\n",
      "\n",
      "Batch size: 32\n",
      "Optimizer: <class 'torch.optim.adam.Adam'>\n",
      "Optimizer params: {'lr': 0.0001}\n",
      "Weight decay: 0.0\n",
      "Max grad norm: None\n",
      "Val dataloader: <torch.utils.data.dataloader.DataLoader object at 0x7f5e0a8b1fd0>\n",
      "Monitor: roc_auc\n",
      "Monitor criterion: max\n",
      "Epochs: 5\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 0 / 5: 100%|██████████| 5/5 [00:00<00:00,  7.38it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-0, step-5 ---\n",
      "loss: 9.4410\n",
      "loss: 9.4410\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 389.26it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-0, step-5 ---\n",
      "roc_auc: 0.0000\n",
      "pr_auc: 0.0476\n",
      "loss: 7.3290\n",
      "New best roc_auc score (0.0000) at epoch-0, step-5\n",
      "\n",
      "roc_auc: 0.0000\n",
      "pr_auc: 0.0476\n",
      "loss: 7.3290\n",
      "New best roc_auc score (0.0000) at epoch-0, step-5\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 1 / 5: 100%|██████████| 5/5 [00:00<00:00, 184.46it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-1, step-10 ---\n",
      "loss: 8.0077\n",
      "loss: 8.0077\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 378.68it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-1, step-10 ---\n",
      "roc_auc: 0.1500\n",
      "pr_auc: 0.0556\n",
      "loss: 4.7464\n",
      "New best roc_auc score (0.1500) at epoch-1, step-10\n",
      "\n",
      "roc_auc: 0.1500\n",
      "pr_auc: 0.0556\n",
      "loss: 4.7464\n",
      "New best roc_auc score (0.1500) at epoch-1, step-10\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 2 / 5: 100%|██████████| 5/5 [00:00<00:00, 185.98it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-2, step-15 ---\n",
      "loss: 4.9490\n",
      "loss: 4.9490\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 454.91it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-2, step-15 ---\n",
      "roc_auc: 0.1500\n",
      "pr_auc: 0.0556\n",
      "loss: 2.3103\n",
      "\n",
      "roc_auc: 0.1500\n",
      "pr_auc: 0.0556\n",
      "loss: 2.3103\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 3 / 5: 100%|██████████| 5/5 [00:00<00:00, 192.87it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-3, step-20 ---\n",
      "loss: 3.4749\n",
      "loss: 3.4749\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 445.30it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-3, step-20 ---\n",
      "roc_auc: 0.2500\n",
      "pr_auc: 0.0625\n",
      "loss: 0.7218\n",
      "New best roc_auc score (0.2500) at epoch-3, step-20\n",
      "\n",
      "roc_auc: 0.2500\n",
      "pr_auc: 0.0625\n",
      "loss: 0.7218\n",
      "New best roc_auc score (0.2500) at epoch-3, step-20\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 4 / 5: 100%|██████████| 5/5 [00:00<00:00, 195.50it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-4, step-25 ---\n",
      "loss: 2.0244\n",
      "loss: 2.0244\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 438.87it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-4, step-25 ---\n",
      "roc_auc: 0.7500\n",
      "pr_auc: 0.1667\n",
      "loss: 0.2822\n",
      "New best roc_auc score (0.7500) at epoch-4, step-25\n",
      "roc_auc: 0.7500\n",
      "pr_auc: 0.1667\n",
      "loss: 0.2822\n",
      "New best roc_auc score (0.7500) at epoch-4, step-25\n",
      "Loaded best model\n",
      "Loaded best model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "trainer.train(\n",
    "    train_dataloader=train_loader,\n",
    "    val_dataloader=val_loader,\n",
    "    epochs=5,\n",
    "    monitor=\"roc_auc\",\n",
    "    optimizer_params={\"lr\": 1e-4},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "264a57fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluation: 100%|██████████| 2/2 [00:00<00:00, 381.40it/s]\n",
      "Evaluation: 100%|██████████| 2/2 [00:00<00:00, 381.40it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'roc_auc': 0.4227642276422764,\n",
       " 'pr_auc': 0.07558781088192854,\n",
       " 'loss': 0.4880794435739517}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.evaluate(test_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "438b7721",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pyhealth",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
